

#from matplotlib.pylab import figure, plot
#matplotlib.interactive(True)
from matplotlib.transforms import Bbox # Point, Value, Interval

from numpy.ma.core import getmaskarray, array as marray, nomask
from numpy import array, logical_or, logical_not, polyfit, poly1d
from numpy import polyfit, poly1d

def maskline(line,bbox,axis=2,outside=False):
    xdata = marray(line.get_xdata(1))
    ydata = marray(line.get_ydata(1))
        
    oldma_x = getmaskarray(xdata)
    oldma_y = getmaskarray(ydata)
    oldma = logical_or(oldma_x,oldma_y)
    #print oldma    
    newma = oldma
    for j,xval in enumerate(xdata):
        yval = ydata[j]
        if axis == 2:
            newma[j] = True if bbox.contains(xval,yval) else newma[j]
        elif axis == 1:
            newma[j] = True if bbox.containsy(yval) else newma[j]
        elif axis == 0:
            newma[j] = True if bbox.containsx(xval) else newma[j]
                    
    #newma = logical_or(oldma,newma)
    
    #print newma
    
    if outside:
        newma = logical_not(newma)
    
    #print newma
    xdata.mask = newma
    ydata.mask = newma
    line.set_xdata(xdata)
    line.set_ydata(ydata)
    #line.recache()
    
    return newma

def maskcollection(coll,bbox,axis=2,outside=False):
    paths = coll.get_paths()
    for path in paths:
        for vert in path.vertices:
            print vert
            if bbox.contains(*vert.tolist()):
                path.vertices = marray(path.vertices,mask=True)
                break
    return coll
  

def mask(axs,bbox,axis=2,outside=False):
    
    for line in axs.lines:
        maskline(line,bbox,axis,outside)
        
        #axs.figure.show()
        #prompt = raw_input("Press any key to continue...")
        
        #newlines.append(line)

    #axs.lines = newlines
    
    for coll in axs.collections:
        maskcollection(coll,bbox,axis,outside)

def unmask(axs,bbox):
#    newlines = []
    for line in axs.lines:
        unmaskline(line,bbox)

def unmaskline(line,bbox):
    xdata = marray(line.get_xdata(0))
    ydata = marray(line.get_ydata(0))
        
    oldma_x = getmaskarray(xdata)
    oldma_y = getmaskarray(ydata)
    oldma = logical_or(oldma_x,oldma_y)
    
    #print oldma
    newma = oldma
    xdata.mask = nomask
    ydata.mask = nomask
    for j,xval in enumerate(xdata):
        yval = ydata[j]
        newma[j] = False if bbox.contains(xval,yval) else newma[j]
    
    #print newma
    xdata.mask = newma
    ydata.mask = newma
    line.set_xdata(xdata)
    line.set_ydata(ydata)
    
    return newma

def fitPolynomial(order,axs,linnum=0):
    line = axs.lines[linnum]

    xdata = line.get_xdata()
    ydata = line.get_ydata()

    xmask = getmaskarray(xdata)
    ymask = getmaskarray(ydata)
    mask = logical_or(xmask, ymask)

    xrange = xdata.compress(-mask)
    yrange = ydata.compress(-mask)

    xrange = xrange.data
    yrange = yrange.data
    
    #print xrange
    #prompt = raw_input("x points...")
    
    #print yrange
    #prompt = raw_input("y points...")
    
    #figure(2)
    #plot(xrange,yrange)
    #figure(1)

    poly = polyfit(xrange,yrange,order)

    return poly

def fitPolynomialBackground(axs,order,x1,x2,trimedge=0,linenum=0):
    
    unmask(axs)
    line = axs.lines[linenum]
    
    p1 = [x1,0] #Point(Value(x1),Value(0))
    p2 = [x2,0] #Point(Value(x2),Value(0))
    bbox = Bbox([p1,p2])
    maskline(line,bbox,axis=0)
    #axs.figure.show()
    #prompt = raw_input("Mask 1...")

    
    xdata = line.get_xdata()
    #print numpy.ma.getmaskarray(xdata)
    minx = min(xdata.data)
    maxx = max(xdata.data)
    #print minx,maxx
    
    p1 = [minx,0] #Point(Value(minx),Value(0))
    p2 = [minx+trimedge,0] #Point(Value(minx+trimedge),Value(0))
    bbox = Bbox([p1,p2])
    maskline(line,bbox,axis=0)
    
    #axs.figure.show()
    #prompt = raw_input("Mask 2...")
    
    p1 = [maxx,0] #Point(Value(maxx),Value(0))
    p2 = [maxx-trimedge,0] #Point(Value(maxx - trimedge),Value(0))
    bbox = Bbox([p1,p2])
    maskline(line,bbox,axis=0)
    
    #axs.figure.show()
    #prompt = raw_input("Mask 3...")
    
    
    poly = fitPolynomial(order,axs,linenum)
    
    return poly


def meanwidth(y,x,x0):
    import scipy
    area = scipy.trapz(y,x)
    mynorm = scipy.interp(x0,x,y)
    return area / mynorm
    
def find_roots(y,x,y0):
    
    n = len(y)
    _roots = []
    for i in range(1,n-1):
        if y[i] == y0:
            _roots += [x[i], ]
        elif y[i] < y0 and y[i+1] > y0:
            _roots += [ (x[i-1] + x[i+1])/2, ]
        elif y[i] > y0 and y[i+1] < y0:
            _roots += [ (x[i-1] + x[i+1])/2, ]
    return _roots
    

def linewidth(y,x,y0,filter=None):
    #x = float(x)
    #y = float(y)
    _root = find_roots(y,x,y0)
    if filter is not None:
        _root = filter(_root)
    else:
        newroot = []
        for val in _root:
            newroot += [ (val,val,val),]
        _root = newroot

    if len(_root) == 2:
        return abs(_root[0][0] - _root[1][0]), (_root[1][1] - _root[0][2]), (_root[1][2] - _root[0][1])
                
    else:
        return -1
    
    